import io
from wand.image import Image
from wand.color import Color
from PyPDF2 import PdfFileReader, PdfFileWriter
import numpy as np
from math import *
from PIL import Image as image
import time
import sys
import os


def convert(filename, res=120):
    """
    将pdf文件转换为同名的png文件
    :param filename:    pdf文件路径
    :param res:
    :return:            生成的同名png文件
    """
    pdf_file = PdfFileReader(open(filename, "rb"))
    pageObj = pdf_file.getPage(0)
    dst_pdf = PdfFileWriter()
    dst_pdf.addPage(pageObj)

    pdf_bytes = io.BytesIO()
    dst_pdf.write(pdf_bytes)
    pdf_bytes.seek(0)

    img_filename = filename.replace('pdf', 'png')
    img = Image(file=pdf_bytes, resolution=res)
    img.format = 'png'
    img.compression_quality = 120
    img.background_color = Color("white")
    img.save(filename=img_filename)
    img.destroy()


def cut_pdf(filename):
    """
    切割pdf白边
    :param filename: pdf文件路径，目前只实现切割第一页
    :return:
    """
    page = PdfFileReader(open(filename, "rb")).getPage(0)
    output = PdfFileWriter()
    x = page.mediaBox.getUpperRight_x()
    y = page.mediaBox.getUpperRight_y()

    # 获得生成的png文件路径
    img_filename = filename.replace('pdf', 'png')
    img = image.open(img_filename)
    width, height = img.size

    img = img.convert('RGBA')

    matrix = np.zeros((width, height))
    for i in range(width):
        for j in range(height):
            (r, g, b, _) = img.getpixel((i, j))
            matrix[i, j] = (r+g+b)/3

    left = 0
    right = 0
    top = 0
    bottom = 0

    # 利用png的像素获得白边的区域
    for i in range(width):
        ind = np.where(matrix[i, :] < 210)[0]
        if len(ind) > 0:
            left = round(i*x/width) - 1
            break

    for i in range(width-1, -1, -1):
        ind = np.where(matrix[i, :] < 210)[0]
        if len(ind) > 0:
            right = ceil(i*x/width) + 1
            break

    for i in range(height):
        ind = np.where(matrix[:, i] < 210)[0]
        if len(ind) > 0:
            top = round(i*y/height) - 1
            break

    for i in range(height-1, -1, -1):
        ind = np.where(matrix[:, i] < 210)[0]
        if len(ind) > 0:
            bottom = ceil(i*y/height) + 1
            break

    # 切割白边
    page.trimBox.lowerLeft = (0, 0)
    page.trimBox.upperRight = (x, y)
    page.cropBox.lowerLeft = (left, y-bottom)
    page.cropBox.upperRight = (right, y-top)
    output.addPage(page)

    # 生成切割后的pdf文件
    output_filename = "./cut_" + str(int(time.time())) +".pdf"
    outputStream = open(output_filename, 'wb')
    output.write(outputStream)
    outputStream.close()

    # 删除中间文件
    os.remove(filename.replace('pdf', 'png'))


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("Error：请正确输入参数 ")
        print("pdfcut 文件路径")
    else:
        convert(sys.argv[1])
        cut_pdf(sys.argv[1])
